 /********************************************************************************************\
 * libeds - extensible data structures library                                                *
 * Copyright (C) 2006, Dan Kruchinin, <just.asgard@gmail.com>                                 *
 *                                                                                            *
 * -> trees/rb_tree.c                                                                         *
 * Red-black tree implementation.                                                             *
 * Red-black tree is a safe-balanced and quite fast[near O(lgN)] binary tree,                 *
 * you can learn more about it here: [http://en.wikipedia.org/wiki/Red-black_tree]            *
 * ----                                                                                       *
 * keep in mind, red-black tree compare function(that must be defined and                     *
 * implemented in userspace) is _obligatory_.                                                 *
 *                                                                                            *
 * This library is free software; you can redistribute it and/or modify it under the terms of *
 * the GNU Lesser General Public License as published by the Free Software Foundation;        *
 * either version 2.1 of the License, or (at your option) any later version.                  *
 *                                                                                            *
 * This library is distributed in the hope that it will be useful, but WITHOUT ANY WARRANTY;  *
 * without even the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  *
 * See the GNU Lesser General Public License for more details.                                *
 *                                                                                            *
 * You should have received a copy of the GNU Lesser General Public License along with this   *
 * library; if not, write to the Free Software Foundation, Inc., 59 Temple Place, Suite 330,  *
 * Boston, MA 02111-1307 USA                                                                  *
 \********************************************************************************************/

#include <stdio.h> /* FIXME: remove after debugging */
#include <stdlib.h>
#include "ds.h"
#include "safe_alloc.h"
#include "errors.h"
#include "trees/bin_tree.h"
#include "trees/rb_tree.h"

static int insert(struct _bit*, void*, void*);
static int delete(struct _bit*, bit_node*);
static bit_node* find(struct _bit*, const void*);
static int node_get_bh(const bit_node*);
static rb_color_t node_get_color(const bit_node*);

static int __join(struct _bit*, bit_node*);
static void __insert_fixup(struct _bit*, bit_node*);
static void __delete_fixup(struct _bit*, bit_node*);
static void __node_set_color(bit_node*, const rb_color_t);
static void __node_set_bh(bit_node*, const int);
static struct _rb_pvt* __alloc_pvt(void);
static void __destroy_all(struct _bit*, bit_node*);
static void __bh_increment(bit_node*);
static void __bh_decrement(bit_node*);

/* this function mustn't return errors */
rb_tree* rb_create(void)
{
  rb_tree* rbt;

  rbt = (rb_tree*)scalloc(1, sizeof((*rbt)));
  _bit_init_main();
  rbt->tree = _bit_allocate();

  /* I use all default binary tree functions "as is" =) */
  BIT_REG_DEFAULT_FUNCS(rbt);
  rbt->insert = insert;
  rbt->delete = delete;
  rbt->find = find;
  /* extra functions: */
  rbt->node_get_color = node_get_color; /* for getting node color */
  rbt->node_get_bh = node_get_bh;       /* and for getting node black height */

  return rbt;
}

rb_tree* rb_join(rb_tree* t1, rb_tree* t2)
{
  rb_tree* nt;
  bit_node* jn;

  /* boring checkings */
  if((t1 == NULL) || (t2 == NULL)) {
	_ds_set_errno(DSERR_EMPTY_DS, DS_ERRPOS_MAIN);
	return NULL;
  }
  if(node_empty(t1->tree->root) || node_empty(t2->tree->root)) {
	_ds_set_errno(DSERR_EMPTY_NODE, DS_ERRPOS_MAIN);
	return NULL;
  }
  if((t1->tree->funcs.cmp == NULL) || (t2->tree->funcs.cmp == NULL)) {
	_ds_set_errno(DSERR_NO_CMP_FUNC, DS_ERRPOS_MAIN);
	return NULL;
  }
  /* two trees should hols keys and values of the same type and must has identical flags setting */
  if((t1->tree->funcs.cmp != t2->tree->funcs.cmp) || (t1->tree->funcs.copy != t2->tree->funcs.copy) ||
	 (t1->tree->funcs.delete != t2->tree->funcs.delete) || (t1->tree->flags != t2->tree->flags)) {
	_ds_set_errno(DSERR_DS_JDIFF, DS_ERRPOS_MAIN);
	return NULL;
  }

  nt = rb_create();
  nt->tree->funcs = t1->tree->funcs;
  nt->tree->flags = t1->tree->flags;

  if(node_get_bh(t1->tree->root) >= node_get_bh(t2->tree->root)) {
	nt->tree->root = t1->tree->root;
	jn = t2->tree->root;
  }
  else {
	nt->tree->root = t2->tree->root;
	jn = t1->tree->root;
  }

  /* walk through the the nt->tree and join join to it nodes of the jn  */
  if(__join(nt->tree, jn) < 0)
	return NULL;

  /* and most interesting moment: destruction! */
  _bit_destroy(t1->tree);
  _bit_destroy(t2->tree);
  sfree(t1);
  sfree(t2);

  return nt;
}

/*
 *
 * split red-black tree on two new,
 * the first is the same as a parent, but without new_root subtree
 * and the second with root - new_root
 *
 */
rb_tree* rb_split(rb_tree* t, bit_node* new_root)
{
  rb_tree* nt;
  bit_node* parent;

  if(t == NULL) {
	_ds_set_errno(DSERR_EMPTY_DS, DS_ERRPOS_MAIN);
	return NULL;
  }
  if(t->tree->funcs.cmp == NULL) {
	_ds_set_errno(DSERR_NO_CMP_FUNC, DS_ERRPOS_MAIN);
	return NULL;
  }
  if(node_empty(new_root) || node_empty(t->root)) {
	_ds_set_errno(DSERR_EMPTY_NODE, DS_ERRPOS_MAIN);
	return NULL;
  }

  nt = rb_create();
  nt->tree->root = new_root;
  nt->tree->funcs = t->tree->funcs;
  nt->tree->flags = t->tree->flags;
  parent = new_root->parent;

  if(!node_empty(parent)) {
	if(parent->left == new_root)
	  parent->left = NULL;
	else
	  parent->right = NULL;
  }
  if(t->tree->root == new_root)
	t->root = NULL;
  
  nt->tree->root->parent = NULL;
  /* new_root must has black color, because red-black tree root is always black */
  __node_set_color(new_root, RBT_CBLACK);

  parent->_nofnodes -= new_root->_nofnodes + 1;
  _bit_recount_nodes(t->tree->root);

  /* and fixing policy is the same as on deleting */
  if(node_get_color(parent) == RBT_CBLACK)
	__delete_fixup(t->tree, parent);

  return nt;
}

void rb_destroy(rb_tree* t)
{
  if(t->tree == NULL)
    return;
  
  __destroy_all(t->tree, t->tree->root);

  _bit_destroy(t->tree);
  sfree(t);
}

static int insert(struct _bit* tree, void* key, void* value)
{
  bit_node *t, *p = NULL;
  int insert = -1; /* this helps to economize some additional checkings */
  ds_cmp_t ret;

  if(tree->funcs.cmp == NULL) {
    _ds_set_errno(DSERR_NO_CMP_FUNC, DS_ERRPOS_MAIN);
    return DSERR_NO_CMP_FUNC;
  }

  t = tree->root;

  /* start "tree walking" */
  while(!node_empty(t)) {
    p = t;
    p->_nofnodes++; /* we add new node, so number of subnodes must be incremented */
	
    if(tree->flags & DS_FNOKEYS) /* if we don't use keys */
      ret = tree->funcs.cmp(t->value, value);
    else
      ret = tree->funcs.cmp(t->key, key);

    switch(ret) {
    case DS_CMP_EQ:
	  /* if our tree should to hold only uniqueue key or values */
      if(tree->flags & DS_FUK) {
		_ds_set_errno(DSERR_DUPLICATE, DS_ERRPOS_MAIN);
		return DSERR_DUPLICATE;
      }
    case DS_CMP_GRT: /* if node [key/value] is less or equal, left subtree will be used */
      t = t->left;
      insert = BIT_LEFT;
      break;
    case DS_CMP_LESS:
      t = t->right;
      insert = BIT_RIGHT;
      break;
    case DS_CMP_ERR:
      _ds_set_errno(DSERR_SYSTEM, DS_ERRPOS_MAIN);
      return DSERR_SYSTEM;
    }
  }
  
  t = _bit_node_new();
  
  /* allocate pvt structure to hold red-black tree specific things */
  if((t->_pvt = __alloc_pvt()) == NULL)
    return _ds_get_errno(DS_ERRPOS_MAIN);
  if(tree->funcs.copy != NULL) {
    if(tree->funcs.copy(t, key, value) == DSERR_SYSTEM) {
      _ds_set_errno(DSERR_SYSTEM, DS_ERRPOS_MAIN);
      return DSERR_SYSTEM;
    }
  }
  /* if copy function doesn't defined, we use simple links copying */
  else {
    t->key = key;
    t->value = value;
  }

  /* take care about node references, color and black height */
  t->parent = p;
  __node_set_color(t, RBT_CRED); /* new node must be red */
  __bh_increment(t);

  if(insert == BIT_LEFT)
    p->left = t;
  else if(insert == BIT_RIGHT)
    p->right = t;
  else
    tree->root = t;

  /*
   * each insertion can breack red-black tree balancin color structure,
   * so it must be fixed to safe a balance
   *
   */
  __insert_fixup(tree, t);
  return 0;
}

/*
 *
 * return node with needed key or value if search was successful
 * return NULL if node with needed key or value wasn't found
 * return BIT_SEARCH_ERR if error occured
 *
 */
static bit_node* find(struct _bit* tree, const void* key)
{
  ds_cmp_t ret;
  bit_node *tmp;
  
  if(tree->funcs.cmp == NULL) {
    _ds_set_errno(DSERR_NO_CMP_FUNC, DS_ERRPOS_MAIN);
    return BIT_SEARCH_ERR;
  }
  if(key == NULL)
	return NULL;

  tmp = tree->root;
  
  while(!node_empty(tmp)) {
    if(tree->flags & DS_FNOKEYS)
      ret = tree->funcs.cmp(tmp->value, key);
    else
      ret = tree->funcs.cmp(tmp->key, key);

    if(ret == DS_CMP_GRT)
      tmp = tmp->left;
    else if(ret == DS_CMP_LESS)
      tmp = tmp->right;
	else if(ret == DS_CMP_ERR) {
	  _ds_set_errno(DSERR_SYSTEM, DS_ERRPOS_MAIN);
	  return BIT_SEARCH_ERR;
	}
    else
      return tmp;
  }

  return NULL;
}

static int delete(struct _bit* tree, bit_node* n)
{
  bit_node *repl, *tmp;

  if(node_empty(n))
	return 0;
  /*
   *
   * if node, that will be deleted doesn't have one of or both childs,
   * one of it child will take deleted node plece in the first case
   * and NULL in the second.
   *
   */
  if(node_empty(n->left) || node_empty(n->right))
    tmp = n;
  /* otherwice its place must be taked by the node, that place before it in sorted order */
  else
    tmp = _bit_node_successor(n);

  /*
   *
   * node deleting can break red-black tree balance, so we need to determine
   * node, that will be used for structure fixing
   *
   */
  if(!node_empty(tmp->left))
    repl = tmp->left;
  else
    repl = tmp->right;
  if(!node_empty(repl))
    repl->parent = tmp->parent;

  /* take care about the references */
  if(!node_empty(tmp->parent)) {
    bit_node* t = tmp->parent;

    while(!node_empty(t)) {
      t->_nofnodes--; /* and take care about subnodes number */
      t = t->parent;
    }
	
    if(tmp->parent->left == tmp)
      tmp->parent->left = repl;
    else
      tmp->parent->right = repl;
  }
  if(tree->root == n) {
    if(tmp != n)
      tree->root = tmp;
    else
      tree->root = repl;
  }

  /* if we have replacement node, we should to make replacement(sorry for tautology) */
  if(tmp != n) {
    tmp->parent = n->parent;

    if(!node_empty(tmp->parent)) {
      if(tmp->parent->left == n)
		tmp->parent->left = tmp;
      else
		tmp->parent->right = tmp;
    }

    tmp->left = n->left;
    tmp->right = n->right;

    if(!node_empty(tmp->left))
      tmp->left->parent = tmp;
    if(!node_empty(tmp->right))
      tmp->right->parent = tmp;

	/* replacement node inherits color and black height from deleted node */
    __node_set_color(tmp, node_get_color(n));
    __node_set_bh(tmp, node_get_bh(n));
  }
  /* tree structure must be fixed only when deleted node has black color, if it red, all is ok */
  if(node_get_color(tmp) == RBT_CBLACK)
    __delete_fixup(tree, repl);
  if(tree->funcs.delete != NULL)
    tree->funcs.delete(n->key, n->value);
  
  /* node killing. mua-ha-ha! */
  sfree(n->_pvt);
  _bit_node_destroy(n);
  _bit_recount_nodes(tmp);

  return 0;
}

static rb_color_t node_get_color(const bit_node* node)
{
  struct _rb_pvt* pvt;
  
  if(node == NULL)
    return RBT_CEMPTY;

  pvt = (struct _rb_pvt*)node->_pvt;
  return pvt->color;
}

static int node_get_bh(const bit_node* node)
{
  struct _rb_pvt* pvt = (struct _rb_pvt*)node->_pvt;
  return pvt->bh;
}

static void __node_set_bh(bit_node* node, const int bh)
{
  struct _rb_pvt* pvt = (struct _rb_pvt*)node->_pvt;
  pvt->bh = bh;
}

/* insert fixing(from Kormen book) */
static void __insert_fixup(struct _bit* tree, bit_node* n)
{
  bit_node* tmp;

  while(node_get_color(n->parent) == RBT_CRED) {
    if(n->parent == n->parent->parent->left) {
      tmp = n->parent->parent->right;

	  /* if uncle is red */
      if(node_get_color(tmp) == RBT_CRED) {
		__node_set_color(n->parent, RBT_CBLACK);
		__node_set_color(tmp, RBT_CBLACK);
		__node_set_color(n->parent->parent, RBT_CRED);
		n = n->parent->parent;
		__bh_increment(n);
      }
      else {
		if(n == n->parent->right) {
		  n = n->parent;
		  _bit_rotate_left(tree, n);
		}
	
		__node_set_color(n->parent, RBT_CBLACK);
		__node_set_color(n->parent->parent, RBT_CRED);
		_bit_rotate_right(tree, n->parent->parent);
      }
    }
    else {
      tmp = n->parent->parent->left;
      
      if(node_get_color(tmp) == RBT_CRED) {
		__node_set_color(n->parent, RBT_CBLACK);
		__node_set_color(tmp, RBT_CBLACK);
		__node_set_color(n->parent->parent, RBT_CRED);
		n = n->parent->parent;
		__bh_increment(n);
      }
      else {
		if(n == n->parent->left) {
		  n = n->parent;
		  _bit_rotate_right(tree, n);
		}
	
		__node_set_color(n->parent, RBT_CBLACK);
		__node_set_color(n->parent->parent, RBT_CRED);
		_bit_rotate_left(tree, n->parent->parent);
      }
    }
  }

  __node_set_color(tree->root, RBT_CBLACK);
}

/* delete fixing(from Kormen book) */
static void __delete_fixup(struct _bit* tree, bit_node* n)
{
  bit_node* tmp;
  rb_color_t color;

  while((n != tree->root) && (node_get_color(n) == RBT_CBLACK)) {
    if(n->parent->left == n) {
      tmp = n->parent->right;
      
      if(node_get_color(tmp) == RBT_CRED) {
		__node_set_color(tmp, RBT_CBLACK);
		__node_set_color(n->parent, RBT_CRED);
		_bit_rotate_left(tree, n->parent);
		tmp = n->parent->right;
      }
      if((node_get_color(tmp->left) == RBT_CBLACK) &&
		 (node_get_color(tmp->right) == RBT_CBLACK)) {
		__node_set_color(tmp, RBT_CRED);
		n = n->parent;
		__bh_decrement(n);
      }
      else {
		if(node_get_color(tmp->right) == RBT_CBLACK) {
		  __node_set_color(tmp->left, RBT_CBLACK);
		  __node_set_color(tmp, RBT_CRED);
		  _bit_rotate_right(tree, tmp);
		  __bh_increment(n->parent);
		  tmp = n->parent->right;
		}
	
		color = node_get_color(n->parent);
		__node_set_color(tmp, color);
		__node_set_color(n->parent, RBT_CBLACK);
		__node_set_color(tmp->right, RBT_CBLACK);
		__node_set_bh(tmp, node_get_bh(n->parent));
		__bh_decrement(n->parent);
		_bit_rotate_left(tree, n->parent);
		n = tree->root;
      }
    }
    else {
      tmp = n->parent->left;
      
      if(node_get_color(tmp) == RBT_CRED) {
		__node_set_color(tmp, RBT_CBLACK);
		__node_set_color(n->parent, RBT_CRED);
		_bit_rotate_right(tree, n->parent);
		tmp = n->parent->left;
      }
      if((node_get_color(tmp->left) == RBT_CBLACK) &&
		 (node_get_color(tmp->right) == RBT_CBLACK)) {
		__node_set_color(tmp, RBT_CRED);
		n = n->parent;
		__bh_decrement(n);
      }
      else {
		if(node_get_color(tmp->left) == RBT_CBLACK) {
		  __node_set_color(tmp->right, RBT_CBLACK);
		  __node_set_color(tmp, RBT_CRED);
		  _bit_rotate_left(tree, tmp);
		  __bh_increment(n->parent);
		  tmp = n->parent->left;
		}
	
		color = node_get_color(n->parent);
		__node_set_color(tmp, color);
		__node_set_color(n->parent, RBT_CBLACK);
		__node_set_color(tmp->left, RBT_CBLACK);
		__node_set_bh(tmp, node_get_bh(n->parent));
		__bh_decrement(n->parent);
		_bit_rotate_right(tree, n->parent);
		n = tree->root;
      }
    }
  }

  __node_set_color(n, RBT_CBLACK);
  tmp = _bit_node_brother(n);

  if(!node_empty(tmp) && (tree->root == n->parent)) {
    if((node_empty(tmp->left) && node_empty(tmp->right)) &&
       (node_empty(n->left) && node_empty(n->right)))
      __bh_decrement(n->parent);
  }
}

static void __destroy_all(struct _bit* tree, bit_node* n)
{
  if(n == NULL)
    return;

  __destroy_all(tree, n->left);
  __destroy_all(tree, n->right);

  if(tree->funcs.delete != NULL)
    tree->funcs.delete(n->key, n->value);
  
  sfree(n->_pvt);
  _bit_node_destroy(n);
}

static void __node_set_color(bit_node* n, const rb_color_t color)
{
  struct _rb_pvt* pvt;

  if(n == NULL)
    return;
  
  pvt = (struct _rb_pvt*)n->_pvt;
  pvt->color = color;
}

static struct _rb_pvt* __alloc_pvt(void)
{
  struct _rb_pvt* pvt;

  pvt = scalloc(1, sizeof(*pvt));
  return pvt;
}

static void __bh_increment(bit_node* n)
{
  struct _rb_pvt* rb_pvt = (struct _rb_pvt*)n->_pvt;
  rb_pvt->bh++;
}

static void __bh_decrement(bit_node* n)
{
  struct _rb_pvt* rb_pvt = (struct _rb_pvt*)n->_pvt;

  if(rb_pvt->bh > 1)
    rb_pvt->bh--;
}

static int __join(struct _bit* tree, bit_node* n)
{
  bit_node *tmp, *parent = NULL;
  int ret, where = -1;
  ds_cmp_t res;
  
  if(node_empty(n))
	return 0;

  if((ret = __join(tree, n->left)) < 0)
	return ret;
  if((ret = __join(tree, n->right)) < 0)
	return ret;

  tmp = tree->root;

  while(!node_empty(tmp)) {
	parent = tmp;

	if(tree->flags & DS_FNOKEYS)
	  res = tree->funcs.cmp(tmp->value, n->value);
	else
	  res = tree->funcs.cmp(tmp->key, n->key);

	switch(res) {
	case DS_CMP_EQ:
	  if(tree->flags & DS_FUK) {
		where = -1;
		tmp = tmp->left;
		break;
	  }
	case DS_CMP_GRT:
	  tmp->_nofnodes++;
	  tmp = tmp->left;
	  where = BIT_LEFT;
	  break;
	case DS_CMP_LESS:
	  tmp->_nofnodes++;
	  tmp = tmp->right;
	  where = BIT_RIGHT;
	  break;
	case DS_CMP_ERR:
	  _ds_set_errno(DSERR_SYSTEM, DS_ERRPOS_MAIN);
	  return DSERR_SYSTEM;
	}
  }

  if(where == BIT_LEFT)
	parent->left = n;
  else if(where == BIT_RIGHT)
	parent->right = n;
  else
	return 0;

  n->parent = parent;
  n->left = n->right = NULL;
  __node_set_bh(n, RBT_CRED);
  n->_nofnodes = 0;
  __insert_fixup(tree, n);

  return 0;
}

